import os
import cv2 as cv
import numpy as np
import json
import shutil
import PIL
from PIL import Image, ImageDraw
from sklearn.cluster import KMeans
import glob
import scipy.misc
import imageio
import skimage.exposure
import skimage.feature
import skimage.filters
import skimage.io
import skimage.morphology


def find_seed(img, n_split=5, sampler_factor=3.0, thresh_pixel=10):
    h, w = img.shape
    pos_l = w // 2
    pos_r = pos_l
    step = int(np.round(w / (n_split * sampler_factor)))

    def worker(img, sampler_pos):
        mask = np.zeros((h, w), dtype=np.uint8)
        mask[:, sampler_pos] = 255
        res = cv.bitwise_and(img, img, mask=mask)
        res_middle = res[:, sampler_pos]
        idx = res_middle != 0
        col_y = np.arange(0, h, dtype=int)
        dots_y = col_y[idx]
        dots_y = dots_y.reshape(-1, 1)
        kmeans = KMeans(n_clusters=3).fit(dots_y)
        centers = kmeans.cluster_centers_
        centers = np.sort(centers, axis=0).ravel()
        for diff in [centers[2] - centers[1], centers[1] - centers[0]]:
            if diff < thresh_pixel:
                raise ValueError(f'diff={diff} < thresh={thresh_pixel}')
        point1 = (sampler_pos, int(np.round((centers[0] + centers[1]) / 2)))
        point2 = (sampler_pos, int(np.round((centers[2] + centers[1]) / 2)))
        return point1, point2

    pos_arr = []
    for i in range(w // step):
        if pos_l >= 0:
            pos_arr.append(pos_l)
        if pos_r <= w - 1:
            pos_arr.append(pos_r)
        pos_l -= step
        pos_r += step
    pos_arr = pos_arr[1:]
    # print(len(pos_arr), pos_arr)  # tmp
    for pos in pos_arr:
        # print(pos)  # tmp
        try:
            p1, p2 = worker(img, pos)
            return p1, p2
            break
        except ValueError as ve:
            pass
    return None, None


def flood_fill(img, point):
    h, w = img.shape
    img3c = np.stack((img, img, img), axis=2)
    mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
    cv.floodFill(img3c, mask, point, (0, 255, 0), 1, 1)
    unit_x = w // 5
    mask_tpl = np.zeros_like(img, dtype=np.uint8)
    area_list = []
    for i in range(5):
        xleft = i * unit_x
        mask = mask_tpl.copy()
        mask[:, xleft:xleft + unit_x] = 255
        res = cv.bitwise_and(img3c, img3c, mask=mask)
        idx = (res == (255, 255, 255))
        res[idx[:, :, 0]] = (0, 0, 0)
        ret, res = cv.threshold(res[:, :, 1], 127, 255, cv.THRESH_BINARY)
        area = cv.countNonZero(res)
        area_list.append(area)
    return area_list

def print_result(output, i, all_areas):
    areas = sum(all_areas, [])

    s = "%s,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d" % (
            i,
            areas[0] + areas[5],
            areas[1] + areas[6],
            areas[2] + areas[7],
            areas[3] + areas[8],
            areas[4] + areas[9],
            areas[0],
            areas[1],
            areas[2],
            areas[3],
            areas[4],
            areas[5],
            areas[6],
            areas[7],
            areas[8],
            areas[9]
            )
    print(s)
    output.write(s)
    output.write("\n")


for path in [
    # "/home/data7/Dataset/Ophthalmology_data/20220509-OK/OK镜",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK3月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK6月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-OK/OK初查",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜3月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜6月",
    # "/home/data7/Dataset/Ophthalmology_data/20220510-软镜/软镜初查",
    "/home/data7/Dataset/Ophthalmology_data/20220511-NewMGD",
    # "/home/data7/Dataset/Ophthalmology_data/20220517-newnormal-fiveequalparts",
    # "/home/data7/Dataset/Ophthalmology_data/正常下睑"
]:
    if not os.path.exists(os.path.join(path, "temp")):
        os.mkdir(os.path.join(path, "temp"))
    if not os.path.exists(os.path.join(path, "chart")):
        os.mkdir(os.path.join(path, "chart"))

    with open("%s/区域.json" % path) as f:
        chart = json.load(f)

    output = open('%s/uas.csv' % path, 'w', encoding='gbk', newline="")
    output.write('文件名,a,b,c,d,e,a1,b1,c1,d1,e1,a2,b2,c2,d2,e2\n')

    for k, v in chart.items():
        fnames = k.split(".")[0]
        # print(fnames)

        try:
            im = Image.open("%s/原图/%s.BMP" % (path, fnames))
            mask = Image.new(im.mode, im.size)
            draw = ImageDraw.Draw(mask)
            for r in v['regions']:
                points = list(zip(r['shape_attributes']['all_points_x'], r['shape_attributes']['all_points_y']))
                draw.polygon(points, outline='white', fill=None)

            im_chart = Image.new(im.mode, im.size)
            draw2 = ImageDraw.Draw(im_chart)
            for r in chart[k]['regions']:
                points = list(zip(r['shape_attributes']['all_points_x'], r['shape_attributes']['all_points_y']))
                draw2.polygon(points, outline='white', fill=None)

            ones = np.nonzero(mask)
            top, bottom = np.min(ones[0]), np.max(ones[0])
            left, right = np.min(ones[1]), np.max(ones[1])

            im = np.array(im)
            im_crop = im[top: bottom, left: right]
            # skimage.io.imsave("%s/images/%s.png" % (path, fnames), im_crop)

            im_chart = np.array(im_chart)
            # print(im_chart)
            chart_crop = im_chart[top: bottom, left: right]
            skimage.io.imsave("%s/temp/%s.png" % (path, fnames), chart_crop)

        except Exception as e:
            print(e, k)
            continue

    files = glob.glob("%s/temp/*.png" % path)
    files = sorted(files)

    for fpath in files:
        basename = os.path.basename(fpath)
        name = basename.split(".")[0]
        img = cv.imread(fpath, cv.IMREAD_GRAYSCALE)
        ret, img = cv.threshold(img, 1, 255, cv.THRESH_BINARY)
        img_1 = img.copy()
        img_2 = img.copy()

        h, w = img.shape
        # cv.line(img, (0, h // 2), (w, h // 2), (255, 0, 0), 2)
        for i in range(5):
            i = i + 1
            num = w // 5
            cv.line(img, (num * i, 0), (num * i, h), (255, 0, 0), 2)

        cv.imwrite("%s/chart/%s" % (path, basename), img)


        points = []
        point1, point2 = find_seed(img_1)
        points.append(point1)
        points.append(point2)

        all_areas = []
        for p in points:
            areas = flood_fill(img_2, p)
            all_areas.append(areas)

        print_result(output, name, all_areas)

    temp_path = os.path.join(path, "temp")
    shutil.rmtree(temp_path)







